"""
HiAgent 文件上传工具
支持上传文件到 HiAgent 服务器
"""
import hashlib
import requests
from typing import Optional


def calculate_sha256(file_bytes: bytes) -> str:
    """计算文件的 SHA256 哈希值"""
    sha256_hash = hashlib.sha256()
    sha256_hash.update(file_bytes)
    return sha256_hash.hexdigest()


class HiAgentFileUploader:
    """HiAgent 文件上传器"""

    def __init__(self, host: str = "agent.swjtu.edu.cn"):
        """
        初始化上传器

        Args:
            host: HiAgent 服务器地址
        """
        self.host = host
        self.mime_types = {
            "pdf": "application/pdf",
            "html": "text/html",
            "csv": "text/csv",
            "xls": "application/vnd.ms-excel",
            "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            "doc": "application/msword",
            "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            "ppt": "application/vnd.ms-powerpoint",
            "pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
            "txt": "text/plain",
            "wav": "audio/wav",
            "jpg": "image/jpeg",
            "jpeg": "image/jpeg",
            "png": "image/png",
            "gif": "image/gif"
        }

    def detect_real_extension(self, file_bytes: bytes) -> Optional[str]:
        """
        检测文件真实格式

        Args:
            file_bytes: 文件字节内容

        Returns:
            文件扩展名，如果无法检测则返回 None
        """
        header = file_bytes[:12]

        if header.startswith(b'\xff\xd8'):
            return 'jpg'
        if header.startswith(b'\x89PNG\r\n\x1a\n'):
            return 'png'
        if header.startswith(b'GIF87a') or header.startswith(b'GIF89a'):
            return 'gif'
        if header.startswith(b'%PDF'):
            return 'pdf'

        return None

    def get_content_type(self, filename: str, file_bytes: bytes) -> tuple[str, str]:
        """
        获取文件的 Content-Type

        Args:
            filename: 文件名
            file_bytes: 文件字节内容

        Returns:
            (content_type, extension) 元组
        """
        # 从文件名获取扩展名
        file_ext = filename.split('.')[-1].lower() if '.' in filename else ''

        # 检测真实格式
        real_ext = self.detect_real_extension(file_bytes)

        # 使用真实格式或文件名扩展名
        final_ext = real_ext or file_ext
        content_type = self.mime_types.get(final_ext, "application/octet-stream")

        return content_type, final_ext

    def upload(self, file_bytes: bytes, filename: str) -> Optional[str]:
        """
        上传文件到 HiAgent 服务器

        Args:
            file_bytes: 文件字节内容
            filename: 文件名

        Returns:
            下载链接 URL，失败返回 None
        """
        try:
            # 计算文件哈希
            sha256_hash = calculate_sha256(file_bytes)
            print(f"[FileUploader] SHA256: {sha256_hash}")

            # 获取文件类型
            content_type, _ = self.get_content_type(filename, file_bytes)
            print(f"[FileUploader] Content-Type: {content_type}")

            # 构造请求头
            headers = {
                "X-Content-Sha256": sha256_hash,
                "Content-Type": content_type,
                "Content-Length": str(len(file_bytes))
            }

            # 构造上传 URL
            url = (
                f"https://{self.host}/api/proxy/up"
                f"?Action=UploadRaw&Version=2022-01-01&Region=cn-north-1"
                f"&Id={sha256_hash}&Expire=2h"
            )
            print(f"[FileUploader] Upload URL: {url}")

            # 发送上传请求
            response = requests.post(url, headers=headers, data=file_bytes, verify=False)
            print(f"[FileUploader] Response status: {response.status_code}")

            response.raise_for_status()

            # 解析响应
            response_data = response.json()
            print(f"[FileUploader] Response data: {response_data}")

            if "Result" in response_data and "Path" in response_data["Result"]:
                path = response_data["Result"]["Path"]
                download_url = (
                    f"https://{self.host}/api/proxy/down"
                    f"?Action=Download&Version=2022-01-01"
                    f"&Path={path}&IsAnonymous=true"
                )
                print(f"[FileUploader] Success! Download URL: {download_url}")
                return download_url

            print(f"[FileUploader] Failed! Response missing Result/Path")
            return None

        except Exception as e:
            print(f"[FileUploader] Exception: {type(e).__name__}: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
